import pandas as pd
import numpy as np
import pymysql
from sklearn.tree import DecisionTreeRegressor
from sklearn.model_selection import train_test_split
from datetime import datetime, timedelta
from sqlalchemy import create_engine, text  # 导入text函数


# 数据库连接配置
db_config = {
    'host': '115tt3208yh53.vicp.fun',
    'port': 13245,
    'user': 'root',
    'password': 'wlw565wlw',
    'database': 'ry-vue',
    'charset': 'utf8mb4'
}


# 1. 从MySQL读取数据
def load_data():
    engine = create_engine(
        f"mysql+pymysql://{db_config['user']}:{db_config['password']}@{db_config['host']}:{db_config['port']}/{db_config['database']}?charset={db_config['charset']}"
    )
    try:
        query = "SELECT * FROM processed_house_data"
        df = pd.read_sql(query, engine)
        df['日期'] = pd.to_datetime(df['日期'])
        region_columns = [col for col in df.columns if col.startswith('区域_')]
        region_names = [col.replace('区域_', '') for col in region_columns]
        print(f"检测到的区域: {region_names}")
        return df, region_columns, region_names
    finally:
        engine.dispose()


# 2. 按区域筛选数据
def filter_data_by_region(df, region_col):
    region_df = df[df[region_col] == 1].copy()
    region_df = region_df.sort_values(by='日期')
    other_regions = [col for col in df.columns if col.startswith('区域_') and col != region_col]
    region_df = region_df.drop(columns=other_regions)
    return region_df


# 3. 数据预处理
def preprocess_data(df):
    drop_cols = ['日期', '平均房价', '价格区间_低', '价格区间_高']
    features = df.drop(columns=drop_cols, errors='ignore')
    target = df['平均房价']
    return features, target


# 4. 训练模型
def train_model(features, target):
    test_size = 0.2 if len(features) > 10 else 0.1
    X_train, X_test, y_train, y_test = train_test_split(
        features, target, test_size=test_size, random_state=42
    )

    model = DecisionTreeRegressor(
        max_depth=6,
        min_samples_split=5,
        random_state=42
    )

    model.fit(X_train, y_train)
    score = model.score(X_test, y_test)
    return model, score


# 5. 生成未来三个月的日期和季度
def generate_future_dates(last_date, months=3):
    future_dates = []
    for i in range(1, months + 1):
        year = last_date.year
        month = last_date.month + i
        if month > 12:
            month -= 12
            year += 1
        future_date = datetime(year, month, 1)
        future_dates.append(future_date)
    future_quarters = [date.strftime('%Y%m') for date in future_dates]
    return future_dates, future_quarters


# 6. 生成未来特征数据
def generate_future_features(df, future_dates, region_col):
    last_row = df.iloc[-1].copy()
    drop_cols = ['日期', '平均房价', '价格区间_低', '价格区间_高']
    feature_cols = [col for col in df.columns if col not in drop_cols]
    last_features = last_row[feature_cols].to_dict()

    future_data = []
    for date in future_dates:
        future_feature = last_features.copy()
        future_feature[region_col] = 1
        future_data.append(future_feature)

    if '环比增长率%' in future_data[0] and len(df) >= 3:
        recent_growth = df['环比增长率%'].tail(3).mean()
        for item in future_data:
            item['环比增长率%'] = recent_growth

    return pd.DataFrame(future_data)


# 7. 预测并保存结果（使用text()包装SQL语句）
def predict_and_save(model, future_features, future_quarters, region_name):
    predictions = model.predict(future_features)
    current_time = datetime.now()

    result = []
    for i in range(len(future_quarters)):
        result.append({
            'area_name': region_name,
            'quarter': future_quarters[i],
            'price_type': 3,
            'avg_price': round(predictions[i], 2),
            'create_time': current_time,
            'update_time': current_time
        })
    result_df = pd.DataFrame(result)

    # 数据库插入修复：用text()包装SQL语句
    engine = create_engine(
        f"mysql+pymysql://{db_config['user']}:{db_config['password']}@{db_config['host']}:{db_config['port']}/{db_config['database']}?charset={db_config['charset']}"
    )
    try:
        with engine.connect() as conn:
            # 核心修复：用text()函数包装SQL字符串
            sql = text("""
                INSERT INTO xian_house_price 
                (area_name, quarter, price_type, avg_price, create_time, update_time)
                VALUES (:area_name, :quarter, :price_type, :avg_price, :create_time, :update_time)
            """)
            # 使用命名参数传递（推荐，更清晰）
            for _, row in result_df.iterrows():
                conn.execute(sql, {
                    'area_name': row['area_name'],
                    'quarter': row['quarter'],
                    'price_type': row['price_type'],
                    'avg_price': row['avg_price'],
                    'create_time': row['create_time'],
                    'update_time': row['update_time']
                })
            conn.commit()
        print(f"成功保存{len(result_df)}条{region_name}的预测结果到xian_house_price表")
    except Exception as e:
        print(f"保存{region_name}预测结果失败: {str(e)}")
    finally:
        engine.dispose()

    return result_df


# 主函数
def main():
    df, region_columns, region_names = load_data()
    print(f"成功加载{len(df)}条数据")

    all_predictions = []

    for region_col, region_name in zip(region_columns, region_names):
        print(f"\n开始处理{region_name}的数据...")
        region_df = filter_data_by_region(df, region_col)
        if len(region_df) < 5:
            print(f"{region_name}数据量不足（{len(region_df)}条），跳过...")
            continue

        features, target = preprocess_data(region_df)
        model, score = train_model(features, target)
        print(f"{region_name}模型R²分数: {score:.4f}")

        last_date = region_df['日期'].max()
        future_dates, future_quarters = generate_future_dates(last_date, months=3)
        print(f"{region_name}未来3个季度: {future_quarters}")

        future_features = generate_future_features(region_df, future_dates, region_col)
        predictions = predict_and_save(model, future_features, future_quarters, region_name)
        all_predictions.append(predictions)
        print(f"{region_name}预测结果:")
        print(predictions[['quarter', 'avg_price']])

    if all_predictions:
        combined = pd.concat(all_predictions, ignore_index=True)
        print("\n所有区域预测结果汇总:")
        print(combined.sort_values(by=['quarter', 'area_name'])[['area_name', 'quarter', 'avg_price']])


if __name__ == "__main__":
    main()